import numpy as np
import math
from imageio import imread
import base

import numpy as np
import math


def genERP(i, j, N):
    val = math.pi/N
    w = math.cos((j - (N/2) + 0.5) * val)
    return w


def compute_map_ws(img):
    """calculate weights for the sphere, the function provide weighting map for a given video
        :img    the input original video
    """
    equ = np.zeros((img.shape[0], img.shape[1]))

    for j in range(0, equ.shape[0]):
        for i in range(0, equ.shape[1]):
            equ[j, i] = genERP(i, j, equ.shape[0])

    return equ


def getGlobalWSMSEValue(mx, my):

    mw = compute_map_ws(mx)

    val = np.sum(np.multiply((mx-my)**2, mw))
    den = val / np.sum(mw)

    return den


def _ws_psnr(image1, image2):
    ws_mse = getGlobalWSMSEValue(image1, image2)

    try:
        ws_psnr = 20 * math.log10(255.0 / math.sqrt(ws_mse))
    except ZeroDivisionError:
        ws_psnr = np.inf

    return ws_psnr


def ws_psnr(img1, img2, mode="RGB"):
    if not img1.shape == img2.shape:
        raise ValueError('输入图片形状不相同')
    assert (mode in ["RGB", "Y"]), "运行模式只能是 RGB 或 Y"

    h, w, c = np.shape(img1)

    if c == 3 and mode == "Y":
        img1, img2 = base.rgb2y(img1, img2)
    return _ws_psnr(img1, img2)
